{-# LANGUAGE TypeOperators, FlexibleContexts, UndecidableInstances #-}

module FreeVarsEx where

import U

import qualified Data.IntSet as IS
import Data.Foldable (Foldable, foldMap)

import Yoko





type Frees = IS.IntSet

bump :: Int -> Frees -> Frees
bump k = IS.map (subtract k) . IS.filter (>= k)





anonFVs :: ULC -> Frees
anonFVs = freeVars



class FreeVars a where freeVars :: a -> Frees

instance FreeVars ULC where
  freeVars (Lam tm) = bump 1 $ freeVars tm
  freeVars (Var i)  = IS.singleton i
  freeVars tm       = freeVars $ ig_from tm

-- through sums
instance FreeVars sum => FreeVars (DCsOf t sum) where
  freeVars = freeVars . unDCsOf
instance (FreeVars a, FreeVars b) => FreeVars (a :+: b) where
  freeVars = foldPlus freeVars freeVars
instance (Generic a, FreeVars (Rep a)) => FreeVars (N a) where
  freeVars = freeVars . rep . unN

-- through products
instance FreeVars U where freeVars = const IS.empty
instance (FreeVars a, FreeVars b) => FreeVars (a :*: b) where
  freeVars = foldTimes IS.union freeVars freeVars

-- through fields
instance FreeVars a => FreeVars (Rec a) where
  freeVars = freeVars . unRec
instance FreeVars (Dep a) where freeVars = const IS.empty
instance (Foldable f, FreeVars a) => FreeVars (Par1 f a) where
  freeVars = foldMap freeVars . unPar1
